package milvus;

import com.alibaba.fastjson.JSONObject;
import io.milvus.client.MilvusServiceClient;
import io.milvus.orm.iterator.QueryIterator;
import io.milvus.param.ConnectParam;
import io.milvus.param.R;
import io.milvus.param.dml.QueryIteratorParam;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.partition.request.CreatePartitionReq;
import io.milvus.v2.service.vector.request.GetReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.QueryReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.response.GetResp;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;

import java.util.*;

/**
 * 单向量搜索：如果您的 Collections 只有一个向量场，请使用 search()方法来查找最相似的实体。
 * 该方法会将您的查询向量与集合中的现有向量进行比较，并返回最匹配的 ID 以及它们之间的距离。作为选项，它还可以返回结果的向量值和元数据。
 */
public class SingleVectorSearch {
    final static String databaseName = "test";
    private static final MilvusClientV2 client;

    final static String CLUSTER_ENDPOINT = "http://c-7d932e12dfccbfff.milvus.aliyuncs.com:19530";

    static {
        // 1. Connect to Milvus server
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(CLUSTER_ENDPOINT)
                .dbName(databaseName)
                .username("root")
                .password("@Milvus520")
                .build();
        client = new MilvusClientV2(connectConfig);
    }

    static final String collectionName = "quick_setup";


    public static void main(String[] args) {
        insert();
        addField();
        singleVectorSearch();
        batchVectorSearch();
        partitionSearch();
        outputFieldSearch();
        filterSearch();
        getByIds();
        getCount();
        queryIterator();
    }

    public static void insert() {
        // 2. Create a collection in quick setup mode
        CreateCollectionReq quickSetupReq = CreateCollectionReq.builder()
                .collectionName(collectionName)
                .dimension(5)
                // 如何测量向量 Embeddings 之间的相似性。
                //可能的值为IP,L2,COSINE,JACCARD,和HAMMING,默认值为已加载索引文件的值
                .metricType(IndexParam.MetricType.IP.name())
                .enableDynamicField(true)
                .build();
        client.createCollection(quickSetupReq);
        GetLoadStateReq loadStateReq = GetLoadStateReq.builder()
                .collectionName(collectionName)
                .build();
        boolean state = client.getLoadState(loadStateReq);
        System.out.println("加载collection:" + state);

        // 3. Insert randomly generated vectors into the collection
        List<String> colors = Arrays.asList(
                "green", "blue", "yellow", "red", "black", "white", "purple", "pink", "orange", "brown", "grey");
        List<JSONObject> dataList = new ArrayList<>();
        for (int i = 0; i < 1000; i++) {
            Random rand = new Random();
            String current_color = colors.get(rand.nextInt(colors.size() - 1));
            int current_tag = rand.nextInt(8999) + 1000;
            JSONObject row = new JSONObject();
            row.put("id", (long) i);
            row.put("vector", Arrays.asList(rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat()));
            row.put("color", current_color);
            row.put("tag", current_tag);
            row.put("color_tag", current_color + "_" + current_tag);
            dataList.add(row);
        }
        InsertReq insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(dataList)
                .build();

        InsertResp insertResp = client.insert(insertReq);
        System.out.println("批量插入数据:" + JSONObject.toJSON(insertResp));


        // 6.1. Create a partition
        CreatePartitionReq partitionReq = CreatePartitionReq.builder()
                .collectionName(collectionName)
                .partitionName("red")
                .build();
        client.createPartition(partitionReq);
        System.out.println("创建red分区");
        partitionReq = CreatePartitionReq.builder()
                .collectionName(collectionName)
                .partitionName("blue")
                .build();
        client.createPartition(partitionReq);
        System.out.println("创建blue分区");

        // 6.2 Insert data into the partition
        dataList = new ArrayList<>();
        for (int i = 1000; i < 1500; i++) {
            Random rand = new Random();
            String current_color = "red";
            int current_tag = rand.nextInt(8999) + 1000;
            JSONObject row = new JSONObject();
            row.put("id", (long) i);
            row.put("vector", Arrays.asList(rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat()));
            row.put("color", current_color);
            row.put("tag", current_tag);
            row.put("color_tag", current_color + "_" + current_tag);
            dataList.add(row);
        }

        insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(dataList)
                .partitionName("red")
                .build();
        insertResp = client.insert(insertReq);
        System.out.println("批量插入(red分区):" + JSONObject.toJSON(insertResp));

        dataList = new ArrayList<>();
        for (int i = 1500; i < 2000; i++) {
            Random rand = new Random();
            String current_color = "blue";
            JSONObject row = new JSONObject();
            row.put("id", (long) i);
            row.put("vector", Arrays.asList(rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat()));
            row.put("color", current_color);
            row.put("color_tag", current_color + "_" + (rand.nextInt(8999) + 1000));
            dataList.add(row);
        }

        insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(dataList)
                .partitionName("blue")
                .build();
        insertResp = client.insert(insertReq);
        System.out.println("批量插入(red分区):" + JSONObject.toJSON(insertResp));
    }


    private static void addField() {
        final String collectionName = "test_add_field";
        CreateCollectionReq quickSetupReq = CreateCollectionReq.builder()
                .collectionName(collectionName)
                .dimension(5)
                // 如何测量向量 Embeddings 之间的相似性。
                //可能的值为IP,L2,COSINE,JACCARD,和HAMMING,默认值为已加载索引文件的值
                .metricType(IndexParam.MetricType.IP.name())
                .enableDynamicField(true)
                .build();
        client.createCollection(quickSetupReq);
        GetLoadStateReq loadStateReq = GetLoadStateReq.builder()
                .collectionName(collectionName)
                .build();
        boolean state = client.getLoadState(loadStateReq);
        System.out.println("加载collection:" + state);


        List<String> colors = Arrays.asList(
                "green", "blue", "yellow", "red", "black", "white", "purple", "pink", "orange", "brown", "grey");
        List<JSONObject> dataList = new ArrayList<>();
        for (int i = 0; i < 10; i++) {
            Random rand = new Random();
            String current_color = colors.get(rand.nextInt(colors.size() - 1));
            int current_tag = rand.nextInt(8999) + 1000;
            JSONObject row = new JSONObject();
            row.put("id", (long) i);
            row.put("vector", Arrays.asList(rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat()));
            row.put("color", current_color);
            row.put("tag", current_tag);
            row.put("color_tag", current_color + "_" + current_tag);
            dataList.add(row);
        }
        InsertReq insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(dataList)
                .build();

        InsertResp insertResp = client.insert(insertReq);
        System.out.println("批量插入数据:" + JSONObject.toJSON(insertResp));


        dataList = new ArrayList<>();
        for (int i = 100; i < 110; i++) {
            Random rand = new Random();
            String current_color = colors.get(rand.nextInt(colors.size() - 1));
            int current_tag = rand.nextInt(8999) + 1000;
            JSONObject row = new JSONObject();
            row.put("id", (long) i);
            row.put("vector", Arrays.asList(rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat(), rand.nextFloat()));
            row.put("color", current_color);
            row.put("name", String.valueOf(i));
            row.put("age", i);
            row.put("tag", current_tag);
            row.put("color_tag", current_color + "_" + current_tag);
            dataList.add(row);
        }
        insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(dataList)
                .build();
        insertResp = client.insert(insertReq);
        System.out.println("批量插入数据:" + JSONObject.toJSON(insertResp));
    }

    public static void singleVectorSearch() {
        // 4. Single vector search
        List<List<Float>> query_vectors = List.of(
                Arrays.asList(0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f));

        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(query_vectors)
                .topK(3) // The number of results to return
                .build();

        SearchResp searchResp = client.search(searchReq);
        System.out.println("单矢量搜索:" + JSONObject.toJSON(searchResp));
    }

    public static void batchVectorSearch() {
        // 5. Batch vector search
        List<List<Float>> query_vectors = Arrays.asList(
                Arrays.asList(0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f),
                Arrays.asList(0.19886812562848388f, 0.06023560599112088f, 0.6976963061752597f, 0.2614474506242501f, 0.838729485096104f)
        );
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(query_vectors)
                .topK(2)
                .build();
        SearchResp searchResp = client.search(searchReq);
        System.out.println("批量向量搜索:" + JSONObject.toJSON(searchResp));
    }

    /**
     * 分区搜索可将搜索范围缩小到集合的特定子集或分区。
     * 这对于数据被分割成逻辑或分类的有组织数据集特别有用，可以通过减少要扫描的数据量来加快搜索操作
     */
    public static void partitionSearch() {
        // 6.3 Search within partitions
        List<List<Float>> query_vectors = List.of(
                Arrays.asList(0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f));
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(query_vectors)
                .partitionNames(List.of("red"))
                .topK(5)
                .build();
        SearchResp searchResp = client.search(searchReq);
        System.out.println("分区搜索(red分区):" + JSONObject.toJSON(searchResp));

        searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(query_vectors)
                .partitionNames(List.of("blue"))
                .topK(5)
                .build();
        searchResp = client.search(searchReq);
        System.out.println("分区搜索(blue分区):" + JSONObject.toJSON(searchResp));
    }

    /**
     * 输出字段搜索
     * 可以在请求中指定output_fields ，以返回包含特定字段的结果
     */
    private static void outputFieldSearch() {
        // 7. Search with output fields
        List<List<Float>> query_vectors = List.of(
                Arrays.asList(0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f));

        SearchReq searchReq = SearchReq.builder()
                .collectionName("quick_setup")
                .data(query_vectors)
                .outputFields(List.of("color", "color_tag"))
                .topK(5)
                .build();
        SearchResp searchResp = client.search(searchReq);
        System.out.println("输出字段搜索:" + JSONObject.toJSON(searchResp));
    }

    public static void filterSearch() {
        // 8. Filtered search
        List<List<Float>> query_vectors = List.of(
                Arrays.asList(0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f));

        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(query_vectors)
                .outputFields(List.of("color", "color_tag"))
                .filter("color_tag like \"red%\"")
                // 范围搜索
                .searchParams(Map.of("radius", 0.1, "range", 1.0))
                .topK(5)
                .build();
        SearchResp searchResp = client.search(searchReq);
        System.out.println("过滤搜索:" + JSONObject.toJSON(searchResp));
    }


    public static void getByIds() {
        // 5. Get entities by ID
        GetReq getReq = GetReq.builder()
                .collectionName(collectionName)
                .ids(Arrays.asList(0L, 1L, 2L))
                // 从分区获取实体
                //.partitionName("red")
                .build();
        GetResp entities = client.get(getReq);
        System.out.println("按ID获取实体:" + JSONObject.toJSON(entities));
    }

    public static void getCount() {
        QueryReq getReq = QueryReq.builder()
                .collectionName(collectionName)
                .filter("")
                .partitionNames(List.of("red"))
                .outputFields(List.of("count(*)"))
                .build();
        QueryResp entities = client.query(getReq);
        System.out.println("计数实体:" + JSONObject.toJSON(entities));
    }

    public static void queryIterator() {
        // 5. Query with iterators
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withUri(CLUSTER_ENDPOINT)
                .withDatabaseName(databaseName)
                .withAuthorization("root", "@Milvus520")
                .build();
        MilvusServiceClient client = new MilvusServiceClient(connectParam);

        QueryIteratorParam queryIteratorParam = QueryIteratorParam.newBuilder()
                .withCollectionName(collectionName)
                .withExpr("color_tag like \"brown_8%\"")
                .withBatchSize(50L)
                .addOutField("id")
                .addOutField("vector")
                .addOutField("color_tag")
                .build();

        R<QueryIterator> queryIteratorRes = client.queryIterator(queryIteratorParam);
        if (queryIteratorRes.getStatus() != R.Status.Success.getCode()) {
            System.err.println(queryIteratorRes.getMessage());
            return;
        }

        QueryIterator queryIterator = queryIteratorRes.getData();
        while (true) {
            List<QueryResultsWrapper.RowRecord> batchResults = queryIterator.next();
            if (batchResults.isEmpty()) {
                queryIterator.close();
                break;
            }

//            List<JSONObject> jsonObject = new ArrayList<>();
            for (QueryResultsWrapper.RowRecord queryResult : batchResults) {
                System.out.println(queryResult);
//                JSONObject row = new JSONObject();
//                row.put("id", queryResult.get("id"));
//                row.put("vector", queryResult.get("vector"));
//                row.put("color_tag", queryResult.get("color_tag"));
//                jsonObject.add(row);
            }
        }
    }
}
